Skip to content

[Feature] Sequence sample unit with exact boundary policies - #4050

Merged
vmoens merged 3 commits into
pytorch:mainfrom
coder-jayp:feat/sequence-unit
Aug 6, 2026
Merged

[Feature] Sequence sample unit with exact boundary policies#4050
vmoens merged 3 commits into
pytorch:mainfrom
coder-jayp:feat/sequence-unit

Conversation

@coder-jayp

Copy link
Copy Markdown
Contributor

Description

Implements the Sequence sampling unit to expand sampled anchors into fixed-length sequences.

Key Features:

  • Robust Boundary Recovery: Leverages utils.find_start_stop_traj to accurately map anchors to their trajectory endpoints, cleanly handling ring buffer seams and storage._last_cursor truncations.
  • Strict Boundary Policies:
    • include_reset: Blindly sweeps forward crossing boundaries.
    • pad: Clamps the sequence at the episode termination and masks trailing steps as invalid.
    • stop: Shifts the anchor backward so the sequence ends exactly on the boundary (falls back to pad if the episode is shorter than the sequence length).
  • Metadata Flow: Expands existing per-anchor info and injects accurate sequence_id, step_in_sequence, and validity_mask for every frame in the B * length output.

Motivation and Context

This change separates the trajectory-range mechanics (sequence generation) from the anchor probability distributions (like Prioritized or Random Sampling), solving the problem where individual samplers previously had to handle both. It enables combinations such as prioritized sequence starts and consistent padding/boundary policies.

Addresses piece 2 of #4039

  • I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

  • New feature (non-breaking change which adds core functionality)

Checklist

  • I have read the CONTRIBUTION guide (required)
  • My change requires a change to the documentation.
  • I have updated the tests accordingly (required for a bug fix or a new feature).
  • I have updated the documentation accordingly.

@pytorch-bot

pytorch-bot Bot commented Jul 25, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4050

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 15 Awaiting Approval

As of commit d15fe7c with merge base b92de7d (image):

AWAITING APPROVAL - The following workflows need approval before CI can run:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 25, 2026
@github-actions github-actions Bot added Feature New feature Documentation Improvements or additions to documentation ReplayBuffers Trainers and removed Feature New feature labels Jul 25, 2026
vmoens pushed a commit to theap06/rl that referenced this pull request Aug 5, 2026
…hor priority semantics

Pieces 3 and 4 (first half) of the pytorch#4039 split, on top of the Sequence
unit from pytorch#4050. The window around each anchor becomes burn_in records
before the anchor, the learning region of length records starting at
it, and bootstrap records after it, with stride spacing the whole
window uniformly. A per-record learning_mask info entry is True
exactly on the learning region. Burn-in never shifts the anchor:
entries before the episode start are invalid and clamp to it;
bootstrap entries obey the episode_boundary policy at episode ends;
defaults reproduce the previous behavior exactly.

Priorities live per anchor: a per-record anchor_index info entry
reports the storage index of each record's sampled anchor (the
original anchor, not the stop-shifted one, since that is what the
sampler's distribution selected), so priorities of sampled sequences
update through the ordinary update_priority path, and per-anchor
sampler entries such as importance weights expand block-constant
across the window. Seeded distribution tests pin that range expansion
does not bias anchor selection for uniform or prioritized sampling.

Part of pytorch#4039.
@vmoens
vmoens force-pushed the feat/sequence-unit branch from 78b3225 to 0825ad5 Compare August 5, 2026 15:13
@github-actions github-actions Bot added the Feature New feature label Aug 5, 2026
@vmoens
vmoens force-pushed the feat/sequence-unit branch from 0825ad5 to 6b1b061 Compare August 5, 2026 16:25
vmoens pushed a commit to theap06/rl that referenced this pull request Aug 5, 2026
…hor priority semantics

Pieces 3 and 4 (first half) of the pytorch#4039 split, on top of the Sequence
unit from pytorch#4050. The window around each anchor becomes burn_in records
before the anchor, the learning region of length records starting at
it, and bootstrap records after it, with stride spacing the whole
window uniformly. A per-record learning_mask info entry is True
exactly on the learning region. Burn-in never shifts the anchor:
entries before the episode start are invalid and clamp to it;
bootstrap entries obey the episode_boundary policy at episode ends;
defaults reproduce the previous behavior exactly.

Priorities live per anchor: a per-record anchor_index info entry
reports the storage index of each record's sampled anchor (the
original anchor, not the stop-shifted one, since that is what the
sampler's distribution selected), so priorities of sampled sequences
update through the ordinary update_priority path, and per-anchor
sampler entries such as importance weights expand block-constant
across the window. Seeded distribution tests pin that range expansion
does not bias anchor selection for uniform or prioritized sampling.

Part of pytorch#4039.
@vmoens

vmoens commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator

Thanks for the contribution — I fixed include_reset for partial/full buffers, non-CPU device handling, storage validation, exports/docs/Hydra configs, and type hints in 6b1b061.

theap06 and others added 3 commits August 6, 2026 14:46
Executable contract for piece 2 of the pytorch#4039 split: Sequence(length,
episode_boundary, done_key) expands each anchor into the length records
that follow it in stored-time order, wrapping ring indices across the
storage seam. Boundary policies: pad keeps the anchor and marks the
tail past an episode end invalid with indices clamped inside the
episode; stop shifts the anchor backward to end exactly at the
boundary, falling back to pad for episodes shorter than length;
include_reset crosses the boundary with all entries valid. The unit
adds per-record sequence_id, step_in_sequence and validity_mask info
entries that surface as TensorDict sample keys, expands per-anchor
sampler entries such as prioritized weights to the record count, and
sample(batch_size=B) returns B*length records. Constructor validates
length and the boundary policy.

Tests are expected to fail until the implementation lands.
…afety, device handling, storage validation, public API

- include_reset now computes against the written length instead of
  max_size: anchors near the write head of a partially filled buffer no
  longer produce out-of-range indices, and on a full ring buffer the
  window clamps at the write cursor instead of splicing the newest data
  with the oldest under an all-True validity mask.
- Keep all index bookkeeping on the sampler's index device and move the
  end flags there before _end_to_start_stop, so non-CPU storages
  (CUDA/MPS) no longer trip cross-device comparisons; returned indices
  live on the same device Transition returns.
- Raise an informative TypeError when the storage is not a
  TensorDict-backed TensorStorage (ListStorage, plain-tensor storages),
  and guard the _last_cursor/_is_full attribute accesses.
- Leave 0-dim info entries untouched instead of crashing on
  repeat_interleave.
- Export Sequence from torchrl.data / torchrl.data.replay_buffers, add
  it to the docs autosummary, and use the public import path in tests.
- Move the utils imports to module top; Literal/NestedKey type hints;
  runnable Examples block; normalize sequence-form done_key to tuple.
- Add TransitionConfig and SequenceConfig Hydra companions with
  registration and cross-references (config/class parity).
- Tests: partial-fill and full-ring include_reset, ListStorage/plain
  tensor errors, scalar info entries, custom nested done_key, non-CPU
  storage device, sample-unit config instantiation.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@vmoens
vmoens force-pushed the feat/sequence-unit branch from 6b1b061 to d15fe7c Compare August 6, 2026 13:47
vmoens pushed a commit to theap06/rl that referenced this pull request Aug 6, 2026
…hor priority semantics

Pieces 3 and 4 (first half) of the pytorch#4039 split, on top of the Sequence
unit from pytorch#4050. The window around each anchor becomes burn_in records
before the anchor, the learning region of length records starting at
it, and bootstrap records after it, with stride spacing the whole
window uniformly. A per-record learning_mask info entry is True
exactly on the learning region. Burn-in never shifts the anchor:
entries before the episode start are invalid and clamp to it;
bootstrap entries obey the episode_boundary policy at episode ends;
defaults reproduce the previous behavior exactly.

Priorities live per anchor: a per-record anchor_index info entry
reports the storage index of each record's sampled anchor (the
original anchor, not the stop-shifted one, since that is what the
sampler's distribution selected), so priorities of sampled sequences
update through the ordinary update_priority path, and per-anchor
sampler entries such as importance weights expand block-constant
across the window. Seeded distribution tests pin that range expansion
does not bias anchor selection for uniform or prioritized sampling.

Part of pytorch#4039.
@vmoens
vmoens merged commit 678520e into pytorch:main Aug 6, 2026
90 of 93 checks passed
vmoens pushed a commit to theap06/rl that referenced this pull request Aug 6, 2026
…hor priority semantics

Pieces 3 and 4 (first half) of the pytorch#4039 split, on top of the Sequence
unit from pytorch#4050. The window around each anchor becomes burn_in records
before the anchor, the learning region of length records starting at
it, and bootstrap records after it, with stride spacing the whole
window uniformly. A per-record learning_mask info entry is True
exactly on the learning region. Burn-in never shifts the anchor:
entries before the episode start are invalid and clamp to it;
bootstrap entries obey the episode_boundary policy at episode ends;
defaults reproduce the previous behavior exactly.

Priorities live per anchor: a per-record anchor_index info entry
reports the storage index of each record's sampled anchor (the
original anchor, not the stop-shifted one, since that is what the
sampler's distribution selected), so priorities of sampled sequences
update through the ordinary update_priority path, and per-anchor
sampler entries such as importance weights expand block-constant
across the window. Seeded distribution tests pin that range expansion
does not bias anchor selection for uniform or prioritized sampling.

Part of pytorch#4039.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. Documentation Improvements or additions to documentation Feature New feature ReplayBuffers Trainers

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants